{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Building a Local MCP Client with LlamaIndex\n",
    "\n",
    "This Jupyter notebook walks you through creating a **local MCP (Model Context Protocol) client** that can chat with a database through tools exposed by an MCP server—completely on your machine. Follow the cells in order for a smooth, self‑contained tutorial."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2  Setup a local LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.ollama import Ollama\n",
    "from llama_index.core import Settings\n",
    "\n",
    "llm = Ollama(model=\"llama3.2\", request_timeout=120.0)\n",
    "Settings.llm = llm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3  Initialize the MCP client and build the agent\n",
    "Point the client at your local MCP server’s **SSE endpoint** (default shown below), and list the available tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.tools.mcp import BasicMCPClient, McpToolSpec\n",
    "\n",
    "mcp_client = BasicMCPClient(\"http://127.0.0.1:8000/sse\")\n",
    "mcp_tools = McpToolSpec(client=mcp_client) # you can also pass list of allowed tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "add_data Add new data to the people table using a SQL INSERT query.\n",
      "\n",
      "    Args:\n",
      "        query (str): SQL INSERT query following this format:\n",
      "            INSERT INTO people (name, age, profession)\n",
      "            VALUES ('John Doe', 30, 'Engineer')\n",
      "        \n",
      "    Schema:\n",
      "        - name: Text field (required)\n",
      "        - age: Integer field (required)\n",
      "        - profession: Text field (required)\n",
      "        Note: 'id' field is auto-generated\n",
      "    \n",
      "    Returns:\n",
      "        bool: True if data was added successfully, False otherwise\n",
      "    \n",
      "    Example:\n",
      "        >>> query = '''\n",
      "        ... INSERT INTO people (name, age, profession)\n",
      "        ... VALUES ('Alice Smith', 25, 'Developer')\n",
      "        ... '''\n",
      "        >>> add_data(query)\n",
      "        True\n",
      "    \n",
      "read_data Read data from the people table using a SQL SELECT query.\n",
      "\n",
      "    Args:\n",
      "        query (str, optional): SQL SELECT query. Defaults to \"SELECT * FROM people\".\n",
      "            Examples:\n",
      "            - \"SELECT * FROM people\"\n",
      "            - \"SELECT name, age FROM people WHERE age > 25\"\n",
      "            - \"SELECT * FROM people ORDER BY age DESC\"\n",
      "    \n",
      "    Returns:\n",
      "        list: List of tuples containing the query results.\n",
      "              For default query, tuple format is (id, name, age, profession)\n",
      "    \n",
      "    Example:\n",
      "        >>> # Read all records\n",
      "        >>> read_data()\n",
      "        [(1, 'John Doe', 30, 'Engineer'), (2, 'Alice Smith', 25, 'Developer')]\n",
      "        \n",
      "        >>> # Read with custom query\n",
      "        >>> read_data(\"SELECT name, profession FROM people WHERE age < 30\")\n",
      "        [('Alice Smith', 'Developer')]\n",
      "    \n"
     ]
    }
   ],
   "source": [
    "tools = await mcp_tools.to_tool_list_async()\n",
    "for tool in tools:\n",
    "    print(tool.metadata.name, tool.metadata.description)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3  Define the system prompt\n",
    "This prompt steers the LLM when it needs to decide how and when to call tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM_PROMPT = \"\"\"\\\n",
    "You are an AI assistant for Tool Calling.\n",
    "\n",
    "Before you help a user, you need to work with tools to interact with Our Database\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4  Helper function: `get_agent()`\n",
    "Creates a `FunctionAgent` wired up with the MCP tool list and your chosen LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.tools.mcp import McpToolSpec\n",
    "from llama_index.core.agent.workflow import FunctionAgent\n",
    "\n",
    "async def get_agent(tools: McpToolSpec):\n",
    "    tools = await tools.to_tool_list_async()\n",
    "    agent = FunctionAgent(\n",
    "        name=\"Agent\",\n",
    "        description=\"An agent that can work with Our Database software.\",\n",
    "        tools=tools,\n",
    "        llm=OpenAI(model=\"gpt-4\"),\n",
    "        system_prompt=SYSTEM_PROMPT,\n",
    "    )\n",
    "    return agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5  Helper function: `handle_user_message()`\n",
    "Streams intermediate tool calls (for transparency) and returns the final response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.agent.workflow import (\n",
    "    FunctionAgent, \n",
    "    ToolCallResult, \n",
    "    ToolCall)\n",
    "\n",
    "from llama_index.core.workflow import Context\n",
    "\n",
    "async def handle_user_message(\n",
    "    message_content: str,\n",
    "    agent: FunctionAgent,\n",
    "    agent_context: Context,\n",
    "    verbose: bool = False,\n",
    "):\n",
    "    handler = agent.run(message_content, ctx=agent_context)\n",
    "    async for event in handler.stream_events():\n",
    "        if verbose and type(event) == ToolCall:\n",
    "            print(f\"Calling tool {event.tool_name} with kwargs {event.tool_kwargs}\")\n",
    "        elif verbose and type(event) == ToolCallResult:\n",
    "            print(f\"Tool {event.tool_name} returned {event.tool_output}\")\n",
    "\n",
    "    response = await handler\n",
    "    return str(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6  Initialize the MCP client and build the agent\n",
    "Point the client at your local MCP server’s **SSE endpoint** (default shown below), build the agent, and setup agent context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.tools.mcp import BasicMCPClient, McpToolSpec\n",
    "\n",
    "\n",
    "mcp_client = BasicMCPClient(\"http://127.0.0.1:8000/sse\")\n",
    "mcp_tool = McpToolSpec(client=mcp_client)\n",
    "\n",
    "# get the agent\n",
    "agent = await get_agent(mcp_tool)\n",
    "\n",
    "# create the agent context\n",
    "agent_context = Context(agent)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "User:  Add to the db: Rafael Nadal whose age is 39 and is a tennis player\n",
      "Calling tool add_data with kwargs {'query': \"INSERT INTO people (name, age, profession) VALUES ('Rafael Nadal', 39, 'Tennis Player')\"}\n",
      "Tool add_data returned meta=None content=[TextContent(type='text', text='true', annotations=None)] isError=False\n",
      "Agent:  The data has been added successfully.\n",
      "User:  fetch data\n",
      "Calling tool read_data with kwargs {'query': 'SELECT * FROM people'}\n",
      "Tool read_data returned meta=None content=[TextContent(type='text', text='1', annotations=None), TextContent(type='text', text='Rafael Nadal', annotations=None), TextContent(type='text', text='39', annotations=None), TextContent(type='text', text='Tennis Player', annotations=None)] isError=False\n",
      "Agent:  Here is the data from the database:\n",
      "\n",
      "1. ID: 1\n",
      "   Name: Rafael Nadal\n",
      "   Age: 39\n",
      "   Profession: Tennis Player\n",
      "User:  add to the db: Roger federer whose age is 42 and is a tennis player\n",
      "Calling tool add_data with kwargs {'query': \"INSERT INTO people (name, age, profession) VALUES ('Roger Federer', 42, 'Tennis Player')\"}\n",
      "Tool add_data returned meta=None content=[TextContent(type='text', text='true', annotations=None)] isError=False\n",
      "Agent:  The data has been added successfully.\n",
      "User:  fetch data\n",
      "Calling tool read_data with kwargs {'query': 'SELECT * FROM people'}\n",
      "Tool read_data returned meta=None content=[TextContent(type='text', text='1', annotations=None), TextContent(type='text', text='Rafael Nadal', annotations=None), TextContent(type='text', text='39', annotations=None), TextContent(type='text', text='Tennis Player', annotations=None), TextContent(type='text', text='2', annotations=None), TextContent(type='text', text='Roger Federer', annotations=None), TextContent(type='text', text='42', annotations=None), TextContent(type='text', text='Tennis Player', annotations=None)] isError=False\n",
      "Agent:  Here is the data from the database:\n",
      "\n",
      "1. ID: 1\n",
      "   Name: Rafael Nadal\n",
      "   Age: 39\n",
      "   Profession: Tennis Player\n",
      "\n",
      "2. ID: 2\n",
      "   Name: Roger Federer\n",
      "   Age: 42\n",
      "   Profession: Tennis Player\n"
     ]
    }
   ],
   "source": [
    "# Run the agent!\n",
    "while True:\n",
    "    user_input = input(\"Enter your message: \")\n",
    "    if user_input == \"exit\":\n",
    "        break\n",
    "    print(\"User: \", user_input)\n",
    "    response = await handle_user_message(user_input, agent, agent_context, verbose=True)\n",
    "    print(\"Agent: \", response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
